import numpy as np
import scipy.io as sio
import os


# folder includes all distortion types of ranking data for tid2013
folder = ['JPEG','GN','GB','JP2K','GNC','MSH','CCL','HFN','IN','ID','BW','NEPN','QN','CA','CS','MGN','CQD']

# 使用绝对路径
dir_rank = os.path.abspath('./') + '/'   # 获取当前目录的绝对路径
base_dir = os.path.dirname(os.path.dirname(dir_rank))  # 获取基础目录

# 确保输出目录存在
os.makedirs(os.path.join(base_dir, 'score_tid2013'), exist_ok=True)

train_file = open(os.path.join(base_dir, 'score_tid2013/tid2013_train.txt'), "w")
val_file = open(os.path.join(base_dir, 'score_tid2013/tid2013_val.txt'), "w")
test_file = open(os.path.join(base_dir, 'score_tid2013/tid2013_test.txt'), "w")

pristine = 'pristine_images'  # The folder of reference images to generate different distortions

# 检查pristine文件夹是否存在
if not os.path.exists(dir_rank + pristine):
    print(f"错误: 找不到原始图像文件夹 '{pristine}'")
    train_file.close()
    val_file.close()
    test_file.close()
    exit(1)

real = os.listdir(dir_rank + pristine)
real.sort()
Num = len(real)
Tr_num = int(Num*0.7)  # 训练集占70%
Val_num = int(Num*0.2)  # 验证集占20%
# 测试集占剩余的10%

FileT_p = []    # To save the image names
scores_p = []   # To save the distortion levels

for i in real:
    # 使用绝对路径
    FileT_p.append(os.path.abspath(dir_rank + pristine + '/' + i))
    scores_p.append(0)

shuff_p = range(Num)
#shuff_p = np.random.permutation(range(Num))         # To decide shuffle the data or not   

for t in folder:
    # 检查当前失真类型文件夹是否存在
    if not os.path.exists(dir_rank + t):
        print(f"警告: 跳过不存在的失真类型文件夹 '{t}'")
        continue
    
    # 写入训练集
    for i in range(Tr_num):
        train_file.write('%s %6.2f\n' % (FileT_p[shuff_p[i]], scores_p[shuff_p[i]]))
    # 写入验证集
    for i in range(Tr_num, Tr_num + Val_num):
        val_file.write('%s %6.2f\n' % (FileT_p[shuff_p[i]], scores_p[shuff_p[i]]))
    # 写入测试集
    for i in range(Tr_num + Val_num, Num):
        test_file.write('%s %6.2f\n' % (FileT_p[shuff_p[i]], scores_p[shuff_p[i]]))  
    
    DisType = os.listdir(t)
    DisType.sort()

    ind = 0
    dis_level = 4      # dis_level +1 = Level of distortion can be chose by changing this variable 
    
    for i in DisType[0:dis_level]:
        # 检查子文件夹是否存在
        if not os.path.exists(dir_rank + t + '/' + i):
            print(f"警告: 跳过不存在的子文件夹 '{t}/{i}'")
            continue
            
        fileN = os.listdir(dir_rank + t+'/'+i)
        fileN.sort()
        ind += 1
        FileT = []
        scores = []
        for j in range(len(fileN)):
            # 使用绝对路径
            FileT.append(os.path.abspath(dir_rank + t + '/' + i +'/'+ fileN[j]))
            scores.append(ind)
        
        # 为每个失真类型创建新的索引排列，确保不会超出范围
        file_count = len(FileT)
        if file_count == 0:
            print(f"警告: 文件夹 '{t}/{i}' 中没有文件，跳过")
            continue  # 如果没有文件，跳过此失真类型
            
        # 确保索引不会超出范围
        shuff = range(file_count)
        #shuff = np.random.permutation(range(file_count))  # 如果需要随机排列
        
        # 计算每个集合的大小
        tr_count = min(int(file_count * 0.7), file_count)
        val_count = min(int(file_count * 0.2), file_count - tr_count)
        test_count = file_count - tr_count - val_count
        
        # 写入训练集
        for i in range(tr_count):
            train_file.write('%s %6.2f\n' % (FileT[shuff[i]], scores[shuff[i]]))
        # 写入验证集
        for i in range(tr_count, tr_count + val_count):
            val_file.write('%s %6.2f\n' % (FileT[shuff[i]], scores[shuff[i]]))
        # 写入测试集
        for i in range(tr_count + val_count, file_count):
            test_file.write('%s %6.2f\n' % (FileT[shuff[i]], scores[shuff[i]]))

train_file.close()
val_file.close()
test_file.close()
print("处理完成！文件已保存到绝对路径")



